from astro.utilities import ismember, indexnear
from astro.io import readtxt
from astro.coord import ang_sep, match
import astro.spec
import gas_gal_hiz.read
import numpy as np
import matplotlib.pyplot as pl
from matplotlib.mlab import rec_append_fields

import xcorr
import xcorr_CIV
from xcorr_CIV import BINEDGES, NBINS, CBINS
from xcorr import wlya, wlyb, Ckms, cosmo, PC


# maximum separation allowed for an LBG from a QSO sight-line
# (degrees).
MAXSEP = 15. / 60.

# not sure what this is
MAXDIST = 50.

def readdata():
    """ Read in qsos, lbgs
    """

    # now only keep lbgs that are in rich's other list
    lbgs = readtxt(
        '/home/nhmc/projects/gas_gal_hiz/lbgs/LBGall_0.5_data_radecz.dat',
        names='num,ra,dec,z')
    #imatch = match(lbgs.ra,lbgs.dec, richlbgs.ra, richlbgs.dec, 2.0)
    const = 1. / (PC * 1.e6)
    losdist = [cosmo.Dc(z) * const  for z in lbgs.z]      # Mpc
    lbgs = rec_append_fields(lbgs, 'losdist', losdist)

    # read QSO ra, dec and emission redshift (note min(zqso) = 2.25)
    qsos0 = gas_gal_hiz.read.qsos('all')

    # find subset of QSOs for cross correlation - all that are within
    # a few arcmin of an LBG.
    c0 = []
    for q in qsos0:
        # separation is in degrees
        sep = ang_sep(q['ra'], q['dec'], lbgs.ra, lbgs.dec)
        minsep = sep.min()
        c0.append(True if minsep < MAXSEP else False)

    c0 = np.array(c0)
    # remove all that only have VIMOS spectra
    #c1 = ~ismember(qsos0.num, num_vimos)
    
    qsos1 = qsos0[c0]
    losdist = [cosmo.Dc(z) * const  for z in qsos1.z]    # Mpc
    qsos1 = rec_append_fields(qsos1, 'losdist', losdist)

    return lbgs, qsos1

def read_random_lbgs():
    filename = ('/home/nhmc/projects/gas_gal_hiz/lbgs/'
                'LBGall_0.5_randoms_radecz.dat')
    rlbgs = readtxt(filename, names='ra,dec,z')
    const = 1. / (PC * 1.e6)
    losdist = [cosmo.Dc(z) * const  for z in rlbgs.z]      # Mpc
    rlbgs = rec_append_fields(rlbgs, ['losdist','num'],
                              [losdist, range(1, len(rlbgs) + 1)])
    return rlbgs
    
def select_close_lbgs(lbgs, ra, dec):
    """ Returns lbgs and separations 
    """ 
    sep = ang_sep(ra, dec, lbgs.ra, lbgs.dec)
    is_close =  sep < MAXSEP
    lbgs1 = lbgs[is_close]
    # if there are no LBGs close enough, skip to the next QSO.
    if len(lbgs1) == 0:
        print 'No LBGs close enough to %s, skipping' % temp
        return [], []

    # angular separation of lbgs from QSO sightline (radians)
    lbg_angsep = sep[is_close] * np.pi / 180.
    # comoving transverse separation
    lbg_angdist = np.array([r*los for r,los in
                            zip(lbg_angsep, lbgs1.losdist)])

    return lbgs1, lbg_angdist


if 0:
    lbgs,qsos = readdata()
    rlbgs = xcorr.read_randomlbgs()

if 1:
    pairIDs_s = [[] for i in range(NBINS)]
    paircounts_s = np.zeros(NBINS, float)
    random_pairIDs_s = [[] for i in range(NBINS)]
    random_paircounts_s = np.zeros(NBINS, float)

    contributing_qsos = []
    for qso in qsos:
        print 'QSO: %s %s' %(qso['fieldqso'], qso['name'])

        ra, dec = qso['ra'], qso['dec']
        
        # lbg selection
        lbgs1, lbg_angdist = select_close_lbgs(lbgs, ra, dec)
        print '%i nearby LBGs' % len(lbgs1)

        ######################################
        # for xi(s)
        ######################################
        pairs, pairids = xcorr_CIV.calc_npairs_s(
            BINEDGES, lbg_angdist, lbgs1.losdist,
            lbgs1.num, [qso['losdist']], [qso['num']], MAXDIST)
        paircounts_s += pairs
        for j,p in enumerate(pairids):
            pairIDs_s[j].extend(p)
         
        # random lbg selection:
        rlbgs1, rlbg_angdist = select_close_lbgs(rlbgs, ra, dec)
        print '%i nearby random LBGs' % len(rlbgs1)

        pairs, pairids = xcorr_CIV.calc_npairs_s(
            BINEDGES, rlbg_angdist, rlbgs1.losdist,
            rlbgs1.num, [qso['losdist']], [qso['num']], MAXDIST)
        random_paircounts_s += pairs
        for j,p in enumerate(pairids):
            random_pairIDs_s[j].extend(p)

# check paircounts_s and random version. Nothing interesting
if 0:
    xcorr_CIV.plot1d(pairIDs_s, random_pairIDs_s, cbins=CBINS, nrandom=20,
           title=r'QSO-LBG $\xi(s)$', prefix='qsolbg_xi_s_logbins')
    idgals = xcorr.unique_lbgids(pairIDs_s)
    idqsos = xcorr.unique_absids(pairIDs_s)

